
from db.mysql import mysql
import asyncio

async def create_or_update_agent(agent_data: dict):
    # 首先检查记录是否存在
    check_query = "SELECT COUNT(*) as count FROM agent WHERE userId = %s"
    result = await mysql.fetch_one(check_query, (agent_data['userId'],))
    
    if result and result['count'] > 0:
        # 如果记录存在，执行更新
        update_query = """
            UPDATE agent SET 
            name = %(name)s, 
            gender = %(gender)s, 
            alias = %(alias)s, 
            placeOfOrigin = %(placeOfOrigin)s, 
            profession = %(profession)s, 
            age = %(age)s, 
            style = %(style)s, 
            other = %(other)s, 
            goal = %(goal)s, 
            template = %(template)s, 
            notes = %(notes)s, 
            updatedAt = NOW()
            WHERE userId = %(userId)s
        """
        await mysql.execute(update_query, agent_data)
    else:
        # 如果记录不存在，执行插入
        insert_query = """
            INSERT INTO agent 
            (userId, name, gender, alias, placeOfOrigin, profession, age, style, other, goal, template, notes, createdAt, updatedAt)
            VALUES (%(userId)s, %(name)s, %(gender)s, %(alias)s, %(placeOfOrigin)s, %(profession)s, %(age)s, %(style)s, %(other)s, %(goal)s, %(template)s, %(notes)s, NOW(), NOW())
        """
        await mysql.execute(insert_query, agent_data)
    
    # 添加短暂延迟确保数据已更新
    await asyncio.sleep(0.1)
    
    # 查询更新后的数据
    updated_agent = await get_agent_by_userid(agent_data['userId'])
    return updated_agent

async def get_agent_by_userid(userId: str):
    query = "SELECT * FROM agent WHERE userId = %s"
    return await mysql.fetch_one(query, (userId,))

async def delete_agent(userId: str):
    query = "DELETE FROM agent WHERE userId = %s"
    result = await mysql.execute(query, (userId,))
    return result